"""
修正边界检查的DQN迷宫实现
"""
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from collections import deque
import random
import matplotlib.pyplot as plt

# 1. 定义迷宫环境 
class MazeEnv:
    def __init__(self):
        self.grid = np.array([
            ['S', ' ', ' ', ' ', ' '],
            [' ', '#', '#', ' ', '#'],
            [' ', '#', ' ', ' ', ' '],
            [' ', '#', '#', '#', ' '],
            [' ', ' ', ' ', ' ', 'G']
        ])
        self.start_pos = (0, 0)
        self.goal_pos = (4, 4)
        self.current_pos = self.start_pos
        self.actions = ['up', 'down', 'left', 'right']  # 动作索引：0,1,2,3
        
    def reset(self):
        self.current_pos = self.start_pos
        return self._pos_to_state(self.current_pos)
    
    def step(self, action):
        x, y = self.current_pos
        new_x, new_y = x, y
        
        # 计算新位置
        if action == 0: new_x = max(0, x - 1)        # 上
        elif action == 1: new_x = min(4, x + 1)      # 下
        elif action == 2: new_y = max(0, y - 1)      # 左
        elif action == 3: new_y = min(4, y + 1)      # 右
        
        # 检查是否撞墙
        if self.grid[new_x, new_y] == '#':
            reward = -1
            done = False
        else:
            self.current_pos = (new_x, new_y)
            if self.current_pos == self.goal_pos:
                reward = 10
                done = True
            else:
                reward = -0.1
                done = False
        
        return self._pos_to_state(self.current_pos), reward, done
    
    def _pos_to_state(self, pos):
        """将坐标位置转换为独热编码状态（25维）"""
        state = np.zeros(25)
        state[pos[0] * 5 + pos[1]] = 1
        return state



# 2. 定义DQN网络
class DQN(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(DQN, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(state_dim, 64),
            nn.ReLU(),
            nn.Linear(64, action_dim)
        )
        
    def forward(self, x):
        return self.fc(x)

# 3. 经验回放缓冲区
class ReplayBuffer:
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)
    
    def push(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))
    
    def sample(self, batch_size):
        batch = random.sample(self.buffer, batch_size)
        states = np.array([x[0] for x in batch])
        actions = np.array([x[1] for x in batch])
        rewards = np.array([x[2] for x in batch])
        next_states = np.array([x[3] for x in batch])
        dones = np.array([x[4] for x in batch])
        return (
            torch.FloatTensor(states),
            torch.LongTensor(actions),
            torch.FloatTensor(rewards),
            torch.FloatTensor(next_states),
            torch.FloatTensor(dones)
        )
    
    def __len__(self):
        return len(self.buffer)

# 4. DQN智能体
class DQNAgent:
    def __init__(self, env):
        self.env = env
        self.state_dim = 25  # 5x5网格的独热编码维度
        self.action_dim = 4  # 上下左右
        
        # 网络和优化器
        self.policy_net = DQN(self.state_dim, self.action_dim) 
        self.target_net = DQN(self.state_dim, self.action_dim) 
        self.target_net.load_state_dict(self.policy_net.state_dict()) # 
        self.optimizer = optim.Adam(self.policy_net.parameters(), lr=0.001)
        
        # 超参数
        self.batch_size = 32 # 批次大小
        self.gamma = 0.9 # 折扣因子
        self.epsilon = 1.0 # 贪婪策略的初始ε
        self.epsilon_min = 0.01 # 贪婪策略的最小值
        self.epsilon_decay = 0.995 # 贪婪策略的衰减率
        self.buffer = ReplayBuffer(1000) # 经验回放缓冲区
    
    def select_action(self, state):
        if random.random() < self.epsilon:
            return random.randint(0, self.action_dim - 1)
        with torch.no_grad():
            state_tensor = torch.FloatTensor(state).unsqueeze(0)
            q_values = self.policy_net(state_tensor)
            return q_values.argmax().item()
    
    def update_model(self):
        if len(self.buffer) < self.batch_size:
            return
        
        # 采样批次
        states, actions, rewards, next_states, dones = self.buffer.sample(self.batch_size)
        
        # 计算目标Q值
        with torch.no_grad():
            next_q = self.target_net(next_states).max(1)[0]
            target_q = rewards + (1 - dones) * self.gamma * next_q
        
        # 计算当前Q值 （这里你可能看不懂，gather是torch的一个python语法，可以自己搜一下。这里表示只取当前动作a的Q值）
        current_q = self.policy_net(states).gather(1, actions.unsqueeze(1))
        
        # 计算损失并更新
        loss = nn.MSELoss()(current_q.squeeze(), target_q)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        
        # 衰减ε
        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)
        
        # 定期同步目标网络
        if random.random() < 0.05:
            self.target_net.load_state_dict(self.policy_net.state_dict())
        
        return loss.item()

# 5. 训练和可视化
def train_and_visualize():
    env = MazeEnv()
    agent = DQNAgent(env)
    rewards = []
    
    for episode in range(300):
        state = env.reset()
        episode_reward = 0
        done = False
        step = 0
        while not done:
            step += 1

            # 执行动作
            action = agent.select_action(state)
            next_state, reward, done = env.step(action)
            agent.buffer.push(state, action, reward, next_state, done)
            state = next_state
            episode_reward += reward
            
            # 更新模型
            agent.update_model()
        
        rewards.append(episode_reward)

        print(f"当前游戏轮次 {episode + 1}, 奖励: {episode_reward:.1f}, ε: {agent.epsilon:.2f}, 该轮游戏的步数: {step}")
    
    # 可视化训练曲线
    plt.plot(rewards)
    plt.xlabel('Episode')
    plt.ylabel('Total Reward')
    plt.title('DQN Training Progress')
    plt.show()
    
    # 可视化最优路径
    state = env.reset()
    path = [env.current_pos]
    done = False
    
    while not done:
        with torch.no_grad():
            state_tensor = torch.FloatTensor(state).unsqueeze(0)
            action = agent.policy_net(state_tensor).argmax().item()
        state, _, done = env.step(action)
        path.append(env.current_pos)
    
    # 绘制迷宫和路径
    grid = env.grid.copy()
    for (x, y) in path[1:-1]:
        grid[x, y] = 'o'
    
    print("\nOptimal Path:")
    for row in grid:
        print(' '.join(row))

if __name__ == "__main__":
    train_and_visualize()